#!/usr/bin/env python
# coding: utf-8

# In[3]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
import seaborn as sns

# 设置Tushare token
ts.set_token('6f4ed3cc824c2d45087664dec52d197dfc411c36f4ed520ed8c35c50')
pro = ts.pro_api()

# 1. 获取两年至少10支股票的日线行情数据
# 这里选择10支股票，时间范围为近两年
stocks = ['000001.SZ', '000002.SZ', '000004.SZ', '000005.SZ', '000006.SZ',
          '000007.SZ', '000008.SZ', '000009.SZ', '000010.SZ', '000011.SZ']
data = []
for stock in stocks:
    df = pro.daily(ts_code=stock, start_date='20230508', end_date='20250508')
    df['ts_code'] = stock
    data.append(df)
df = pd.concat(data, ignore_index=True)

# 2. 对股票行情数据进行分类（打标签）与技术指标计算（不少于15个）
# 这里简单定义一个分类：上涨为1，下跌为0
df['label'] = np.where(df['pct_chg'] > 0, 1, 0)

# 计算技术指标
df['MA5'] = df['close'].rolling(window=5).mean()
df['MA10'] = df['close'].rolling(window=10).mean()
df['MA20'] = df['close'].rolling(window=20).mean()
df['MACD'] = df['close'].ewm(span=12).mean() - df['close'].ewm(span=26).mean()
df['RSI'] = df['pct_chg'].rolling(window=14).mean()
df['BollingerUpper'] = df['close'].rolling(window=20).mean() + 2 * df['close'].rolling(window=20).std()
df['BollingerLower'] = df['close'].rolling(window=20).mean() - 2 * df['close'].rolling(window=20).std()
df['CCI'] = (df['close'] - df['close'].rolling(window=14).mean()) / (0.015 * df['close'].rolling(window=14).std())
df['OBV'] = np.where(df['pct_chg'] > 0, df['vol'], -df['vol']).cumsum()

# 计算ATR
df['TR1'] = df['high'] - df['low']
df['TR2'] = (df['high'] - df['close'].shift(1)).abs()
df['TR3'] = (df['low'] - df['close'].shift(1)).abs()
df['TR'] = df[['TR1', 'TR2', 'TR3']].max(axis=1)
df['ATR'] = df['TR'].rolling(window=14).mean()

df['ROC'] = df['close'].pct_change(10)
df['MFI'] = df['vol'] * df['close']
# 修正PSY指标计算
df['PSY'] = pd.Series(np.where(df['pct_chg'] > 0, 1, 0)).rolling(window=12).mean()
df['ADX'] = 0  # 简化示例，这里不详细计算ADX
df['VIX'] = df['pct_chg'].rolling(window=30).std()

# 3. 对分类计算后的数据进行建模前的处理与分析
# 处理空值
df = df.dropna()

# 归一化
scaler = MinMaxScaler()
features = df[['MA5', 'MA10', 'MA20', 'MACD', 'RSI', 'BollingerUpper', 'BollingerLower',
               'CCI', 'OBV', 'ATR', 'ROC', 'MFI', 'PSY', 'ADX', 'VIX']]
features = scaler.fit_transform(features)

# 主成分分析
pca = PCA(n_components=0.95)
features = pca.fit_transform(features)

# 相关性分析
corr = pd.DataFrame(features, columns=[f'PC{i+1}' for i in range(features.shape[1])]).corr()
sns.heatmap(corr, annot=True, cmap='coolwarm')
plt.title('Principal Components Correlation')
plt.show()

# 个股画像（选择6个属性）
stock_profiles = df[['ts_code', 'MA5', 'MA10', 'MA20', 'MACD', 'RSI', 'CCI']].groupby('ts_code').mean()
print("个股画像：")
print(stock_profiles)

# 数据均衡
from imblearn.under_sampling import RandomUnderSampler
X = features
y = df['label']
rus = RandomUnderSampler(random_state=42)
X_resampled, y_resampled = rus.fit_resample(X, y)

# 4. 使用一种机器学习方法分析建模与模型评价，并对评价结果进行可视化
# 使用逻辑回归进行建模
X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=0.2, random_state=42)
model = LogisticRegression()
model.fit(X_train, y_train)
y_pred = model.predict(X_test)

# 模型评价
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率：{accuracy}")
print("分类报告：")
print(classification_report(y_test, y_pred))
confusion = confusion_matrix(y_test, y_pred)
sns.heatmap(confusion, annot=True, fmt='d', cmap='Blues')
plt.xlabel('预测值')
plt.ylabel('真实值')
plt.title('混淆矩阵')
plt.show()


# In[2]:


pip install -i https://pypi.tuna.tsinghua.edu.cn/simple imbalanced-learn


# In[ ]:




